package top.kangyaocoding.tech.controller;

import io.swagger.v3.oas.annotations.Operation;
import jodd.util.StringUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import top.kangyaocoding.tech.model.ChatSession;
import top.kangyaocoding.tech.service.IChatSessionService;

import java.util.UUID;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * 描述: 聊天绘画记忆控制类
 *
 * @author K·Herbert
 * @since 2025-04-13 23:16
 */
@Slf4j
@RestController
@RequestMapping("/api/v1/ai")
@CrossOrigin(origins = "*")
public class ChatControllerWithMemory {

    private final ChatClient chatClient;
    private final IChatSessionService chatSessionService;

    public ChatControllerWithMemory(ChatClient chatClient, IChatSessionService chatSessionService) {
        this.chatClient = chatClient;
        this.chatSessionService = chatSessionService;
    }

    @Operation(summary = "流式回答聊天")
    @GetMapping(value = "generateStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "你是谁") String message, @RequestParam String sessionId, @RequestParam String userId) {
        // 默认生成一个会话
        if (StringUtil.isBlank(sessionId)) {
            sessionId = UUID.randomUUID().toString();
            ChatSession chatSession = ChatSession.builder()
                    .sessionId(sessionId)
                    .sessionName(message.length() >= 15 ? message.substring(0, 15) : message)
                    .build();
            chatSessionService.saveSession(chatSession, userId);
        }
        String finalSessionId = sessionId;
        return chatClient.prompt().user(message)
                .advisors(advisorSpec -> advisorSpec
                        .param(CHAT_MEMORY_CONVERSATION_ID_KEY, finalSessionId)
                        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .stream()
                .chatResponse();
    }

}
